"""Set of classes to implement a CTPRPC Client

You can quickly implement a client by simply instancing :class:`CTPRPCClient`

>>> cli = CTPRPCClient(('localhost', 1110))
>>> print cli.echo()
HELO
>>> 
"""

import socket, select, os, threading, time, traceback, sys

from threading import Thread
from Queue import Queue

from CTPRPC3Common import *
from CTPRPCPickle import  CTPRPCSerializer

try:
    from tlslite.api import *
except:
    print 'System wide tlslite not found, falling back to internal one'
    from os3.net.externals.tlslite.api import *

try:
    from hashlib import md5
except:
    from md5 import md5

class CTPRPCCallTracker(object):
    def __init__(self, serializer):
        self.serializer = serializer

        self.requested = False
        self.sent = False
        self.timestamp = 0
        self.call_data = None

    def mark_in_progress(self):
        if self.sent:
            self.update_request_time()

    @property
    def has_timed_out(self):
        return time.time() - self.timestamp > 6

    def update_request_time(self):
        self.timestamp = time.time()

    def got_answer(self, ans):
        if ans['id'] >= self.call_data['id']:
            self.requested = False
            self.sent = False
            self.timestamp = 0
            self.call_data = None
            return True

        return False

    def new_call(self, call_data):
        self.requested = True
        self.sent = False
        self.timestamp = time.time()
        self.call_data = call_data['data']
        self.sent = False

    @property
    def call_request(self):
        return CallPacket(self.serializer, self.call_data)

    def mark_as_sent(self):
        self.sent = True
        self.update_request_time()

class CTPRPCClientAuthenticator(object):
    """
    Will perform remote side connection and authentication for a :class:`CTPRPCClient`.
    This will be instanced by the :class:`CTPRPCClient` and you won't need to use this directly
    
    It will handle the complete handshaking session by using **tlslite** for encryption
    and authentication.
    """
    def __init__(self, server, username=None, password=None):
        self.user = username
        if password:
            self.password = md5(password).hexdigest()
        else:
            self.password = None
        self.server = server

        self.current_session = None
        self.current_dest = None

    def _parse_handshake(self, what):
        try:
            self.current_session, p = map(lambda x : int(x), what.split(':'))
            self.current_dest = (self.server[0], p)
            return (self.current_session, self.current_dest)
        except Exception, e:
            raise CTPRPCError('Bad handshake response from server')

    def connect(self):
        """
        Will connect  to the remote server and perform authentication.
        
        After connecting it will return a **SESSION_ID** and the remote server **(ADDRESS, UDP PORT)**
        """
        sock = socket.socket()
        sock.connect(self.server)
        sock.settimeout(10)

        if self.user:
            connection = TLSConnection(sock)
            connection.closeSocket = True
            connection.handshakeClientSRP(self.user, self.password)
        else:
            connection = sock

        what = connection.recv(4096)
        connection.close()

        return self._parse_handshake(what)

class CTPRPCClientDataDetector(Thread):
    def __init__(self, queue, sock):
        super(CTPRPCClientDataDetector, self).__init__()
        self.io_queue = queue
        self.sock = sock
        self.keep_running = True

    def run(self):
        while self.keep_running:
            try:
                rl, wl, xl = self.select ([self.sock], [], [], 5)
            except:
                break

            if self.keep_running:
                if not xl and not rl and not wl:
                    self.io_queue.put(queue_msg('HBEAT'))

                if self.sock in rl:
                    self.io_queue.put(queue_msg('NET'))

                self.io_queue.join()

    def stop(self):
        """Stops the data detector, Call this method ONLY FROM ANOTHER THREAD"""
        self.keep_running = False
        if self.isAlive():
            self.join(1.0)
            try:
                self.io_queue.get(timeout=1.0)
            except:
                pass

class CTPRPCClientIO(Thread):
    """
    IO Layer of the Client.
    This will be instanced by the :class:`CTPRPCClient` and you won't need to use this directly
    
    It will receive commands from the :class:`CTPRPCClient`, transforms them in packets and send
    to the server. It will also receive any response or event from the server handling them or
    forwarding them to the :class:`CTPRPCClient` instance.
    
    IO Layer will also handle timeouts, automatic reconnections, heartbeats and events.
    """
    def __init__(self, ipc, authenticator, parent_thread=None):
        super(CTPRPCClientIO, self).__init__()
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.auth = authenticator
        self.keep_running = True
        self.parent_thread = parent_thread
        self.serializer = CTPRPCSerializer()

        self.master_queue = ipc[0]
        self.io_queue = ipc[1]

        self.session = self.auth.current_session
        self.dest = self.auth.current_dest
        self.seq = 0
        self.recv_seq = 0
        self.call_tracker = CTPRPCCallTracker(self.serializer)
        self.data_detector = CTPRPCClientDataDetector(self.io_queue, self.sock)

        self.events = {}

        #Those are to permit monkey patching of the network layer
        self.select = select.select
        self.current_packet = None

    def send(self, pkt):
        """
        Sends the specified :class:`os3.net.CTPRPC3Common.CTPRPCPacket` instance.
        Before sending it the packet will be splitted in multiple
        packets that can be contained inside the MTU
        """
        if not self.sock or not self.keep_running:
            return None

        pkt.seq = self.seq + 1
        pkt.session = self.session

        for pkt in pkt.generated_split():
            self.seq = pkt.seq
            if not reliable_send(self, pkt, self.dest):
                self.master_queue.put(queue_msg('WARNING', 'Reliable data send failed, will repost the current request from begin'))
                return

    def parse_master_request(self, msg):
        """
        Will handle the current :class:`CTPRPCClient` request
        pending inside the communication queue
        """

        #If the message type is a method call
        if msg['type'] == 'CALL':
            try:
                self.call_tracker.new_call(msg)
                self.send(self.call_tracker.call_request)
                self.call_tracker.mark_as_sent()
                self.master_queue.put(queue_msg('OK'))
            except Exception, e:
                traceback.print_exception(*sys.exc_info())
                self.master_queue.put(queue_msg('ERROR', str(e)))
        elif msg['type'] == 'SUBSCRIBE':
            event_name = msg['data']['event_name']
            callback = msg['data']['callback']

            if not self.events.has_key(event_name):
                self.events[event_name] = []
                pkt = SubscribePacket(event_name)
                self.send(pkt)

            self.events[event_name].append(callback)
        elif msg['type'] == 'UNSUBSCRIBE':
            event_name = msg['data']['event_name']
            callback = msg['data']['callback']

            if self.events.has_key(event_name):
                self.events[event_name].remove(callback)

        elif msg['type'] == 'QUIT':
            self.stop()
            self.master_queue.put(queue_msg('QUIT'))
        else:
            self.master_queue.put(queue_msg('ERROR'))

    def parse_incoming_packet(self):
        """
        Will handle the currently available packet sent by the server
        """
        try:
            data, addr = self.sock.recvfrom(1500, socket.MSG_DONTWAIT)
        except:
            return

        try:
            pkt = CTPRPCPacket.load(data)
        except Exception, e:
            self.current_packet = None
            self.master_queue.put(queue_msg('WARNING', 'Corrupted Packet, Packet discarded'))
            return

        self.sock.sendto('K-%s' % pkt.seq, self.dest)
        if self.recv_seq < pkt.seq:
            self.recv_seq = pkt.seq
        else:
            return

        try:
            if self.current_packet and not self.current_packet.marker:
                self.current_packet.append(pkt)
            else:
                self.current_packet = pkt

            if not self.current_packet.marker:
                return

            if not self.current_packet.is_valid:
                self.current_packet = None
                self.master_queue.put(queue_msg('WARNING', 'Invalid CRC, Packet discarded'))
                return

            self.current_packet.decompress()
            self.current_packet = parse_packet(self.current_packet)

            if isinstance(self.current_packet, HBeatPacket):
                if not self.call_tracker.sent:
                    return

                if self.current_packet.status == 'C':
                    self.call_tracker.mark_in_progress()
                else:
                    if self.call_tracker.has_timed_out:
                        self.master_queue.put(queue_msg('WARNING', 'Current method Call timeout, retrying...'))
                        self.send(self.call_tracker.call_request)
                        self.call_tracker.update_request_time()
            elif isinstance(self.current_packet, CallPacket):
                try:
                    p = self.current_packet.get_parsed_method(self.serializer)
                except:
                    return

                event_name = p['method']
                if self.events.has_key(event_name):
                    for call in self.events[event_name]:
                        try:
                            call(*p['args'], **p['kw'])
                        except Exception, e:
                            traceback.print_exc()
            else:
                if not self.call_tracker.sent:
                    return

                try:
                    ans = self.current_packet.get_parsed_result(self.serializer)
                except:
                    self.master_queue.put(queue_msg('ERROR', 'Data deserialization failed even though packet was valid'))
                    traceback.print_exc()
                    return

                if self.call_tracker.got_answer(ans):
                    self.master_queue.put(queue_msg('PACKET', self.current_packet.payload))
        except Exception, e:
            traceback.print_exception(*sys.exc_info())
            self.master_queue.put(queue_msg('ERROR', str(e)))

    def run(self):
        """
        This will simply loop handling network packets and :class:`CTPRPCClient` requests.
        it is meant to be called by the threading module using the :meth:`start` method.
        """
        hbeat_sent = 0
        self.data_detector.select = self.select
        self.data_detector.start()

        while self.keep_running:
            if self.parent_thread and not self.parent_thread.isAlive():
                self.stop()
                break

            got_packet = False
            try:
                msg = self.io_queue.get(timeout=7)
                got_packet = True
            except:
                msg = {'type': 'HBEAT'}

            if msg['type'] not in ('NET', 'HBEAT'):
                self.parse_master_request(msg)
            else:
                if self.keep_running:
                    if msg['type'] == 'HBEAT':
                        hbeat_sent += 1
                        self.send(HBeatPacket())

                    if msg['type'] == 'NET':
                        hbeat_sent = 0
                        self.parse_incoming_packet()

                    if hbeat_sent > 2:
                        hbeat_sent = 0
                        self.keep_running = self.reconnect()
                        if self.keep_running:
                            self.master_queue.put(queue_msg('RECONNECTED'))
                        else:
                            self.master_queue.put(queue_msg('TIMEOUT'))
                            self.stop()

            if got_packet:
                self.io_queue.task_done()

    def stop(self):
        """
        Will stop the IO layer from handling any other packet or :class:`CTPRPCClient` request. 
        """
        self.keep_running = False

        self.data_detector.stop()
        self.sock.close()

    def reconnect(self):
        """
        Will force the IO layer to perform a reconnection.
        
        If the server just completed execution of a method and you force
        reconnection before receiving the method AckPacket you will cause
        the client to execute the method again.
        """
        try:
            self.session, self.dest = self.auth.connect()
            self.seq = 0

            self.data_detector.stop()
            self.sock.close()

            self.sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            self.data_detector = CTPRPCClientDataDetector(self.io_queue, self.sock)
            self.data_detector.select = self.select
            self.data_detector.start()
            return True
        except:
            return False

class CTPRPCClient(object):
    """
    :class:`CTPRPCClient` is the CTPRPC interface to the server.
    
    It permits to connect to a remote CTPRPC **server** in the form *(host, post)*, 
    call rpc methods and to subscribe/unsubscribe to remote events.
    
    If you give a **user** and **pwd** it will authenticate to the remote side using encryption,
    authentication is handled by a :class:`CTPRPCClientAuthenticator`
    
    It performs handshaking and then forks an instance of a :class:`CTPRPCClientIO`
    to handle all the network layer. It comunicates with that instance using
    two shared queues to pass data.

    You can prevent it from automatically connect to the remote server by providing **delay_connect** parameter
    """
    def __init__(self, server, user=None, pwd=None, delay_connect=False):
        self.server = server
        self.logger = CTPRPCClientEventsLogger()
        self.authenticator = CTPRPCClientAuthenticator(server, user, pwd)
        self.serializer = CTPRPCSerializer()
        self.call_lock = threading.Lock()

        self.io = None
        self.session = None
        self.last_call_request_id = 0
        self.last_excp = None

        if not delay_connect:
            self.connect()

    def quit(self):
        """
        Closes the client communication and quits the :class:`CTPRPCClientIO` instance.
        
        *You should not need to call this as the* :class:`CTPRPCClientIO` *instance will correctly
        quit when its parent* :class:`CTPRPCClient` *instance is destroyed*
        """
        if self.io:
            self.io_queue.put(queue_msg('QUIT', None))

    def connect(self):
        """
        Connects to the server specified inside constructor, will work only if you are
        not already connected. Reconnections will be handled automatically by the
        :class:`CTPRPCClientIO` instance
        """
        if not self.io:
            try:
                self.authenticator.connect()
            except Exception, e:
                self.logger.err_authentication()
                raise e

            self.master_queue = Queue()
            self.io_queue = Queue()

            self.io = CTPRPCClientIO((self.master_queue, self.io_queue), self.authenticator, threading.currentThread())
            self.io.start()

    def _recover_files(self, ans):
        for file in self.serializer.files:
            filename, fileo = file
            fileo.write(ans['files'][filename])
            fileo.seek(0)
        self.serializer.files = []

    def event_subscribe(self, event_name, callback):
        """
        Subscribe for notification to a remote event, when the server will raise the **event_name** event
        **callback** function will be called.
        
        Each event will require its own documentation for callback parameters,
        *multiple callbacks can be registered for one event*
        """
        subscription = {'event_name':event_name,
                        'callback':callback}
        self.io_queue.put(queue_msg('SUBSCRIBE', subscription))

    def event_unsubscribe(self, event_name, callback):
        """
        Unsubscribe the given **callback** from **event_name** event
        """
        subscription = {'event_name':event_name,
                        'callback':callback}
        self.io_queue.put(queue_msg('UNSUBSCRIBE', subscription))

    def call_method(self, method, args, kw):
        """
        Calls the given remote **method** passing **args** list of unamed
         arguments and **kw** dictionary of named arguments.
         
        *You probably won't need to use this method directly as you can
        simply call the remote method by performing instance.method(*args, **kw)*
        
        If the execution of the method raises any exception this will be reraised as
        a :class:`os3.net.CTPRPC3Common.RemoteException` instance
        """
        self.call_lock.acquire()

        call_request = {'id':self.last_call_request_id,
                        'method':method,
                        'args':args,
                        'kw':kw}
        self.last_call_request_id += 1
        self.io_queue.put(queue_msg('CALL', call_request))

        got_answer = False
        ans = None
        while not got_answer:
            answer = self.master_queue.get()
            if answer['type'] == 'RECONNECTED':
                self.logger.reconnected(answer['time'])
            elif answer['type'] == 'WARNING':
                self.logger.warning(answer['time'], answer['data'])
            elif answer['type'] == 'TIMEOUT':
                if not self.logger.err_timeout(answer['time']):
                    self.call_lock.release()
                    raise CTPRPCError('Timeout')
            elif answer['type'] == 'ERROR':
                if not self.logger.err_protocol_io(answer['time'], answer['data']):
                    self.call_lock.release()
                    raise BadMethodError('IO Layer Error while calling %s: %s' % (method, answer['data']))
            elif answer['type'] == 'OK':
                self.logger.sent_call(answer['time'], method)
            elif answer['type'] == 'QUIT':
                self.logger.quit(answer['time'], method)
            elif answer['type'] == 'PACKET':
                self.logger.got_call_answer(answer['time'], method, answer['data'])

                pkt = parse_packet(CTPRPCPacket(answer['data']))
                if isinstance(pkt, AckPacket):
                    ans = pkt.get_parsed_result(self.serializer)
                    self._recover_files(ans)
                    ans = ans['res']
                elif isinstance(pkt, ExceptionPacket):
                    ans = pkt.get_parsed_result(self.serializer)
                    if not self.logger.err_remote_exception(answer['time'], method, ans):
                        self.last_excp = ans['excp']
                        self.call_lock.release()
                        raise ans['excp']
                    ans = ans['excp']
                elif isinstance(pkt, NotFoundPacket):
                    if not self.logger.err_method_not_found(answer['time'], method):
                        self.call_lock.release()
                        raise BadMethodError('method "%s" not found on remote object' % method)
                got_answer = True

        self.call_lock.release()
        return ans

    def __getattr__(self, attr):
        class CTPRPCMethod (object):
            def __init__ (self, client, method_name):
                super (CTPRPCMethod, self).__init__()
                self.client = client
                self.name = method_name

            def __getattr__ (self, name):
                return CTPRPCMethod (self.client, '%s.%s' % (self.name, name))

            def __call__ (self, *args, **kw):
                return self.client.call_method(self.name, args, kw)

        return CTPRPCMethod(self, attr)

def destroy_client(cli):
    """destroy_client has been deprecated and is here only for compatibility with old api.
    New client class will destroy its instances when the garbage collector retains them."""
    print "CTPRPC Warning: desploy_client has been deprecated and isn't required anymore"

def set_close_handler(cli, handler):
    """set_close_handler has been deprecated and is here only for compatibility with old api.
    New client class will close the socket exactly when the object is deleted, so
    it won't need anymore a callback for that event"""
    print "CTPRPC Warning: set_close_handler has been deprecated as the socket will be automatically closed when destroying the client"

def enable_encryption(client, enabled=1):
    """enable_encryption has been deprecated and is here only for compatibility with old api.
    New client class will use encryption automatically if connection has been instanced with a username and password"""
    print "CTPRPC Warning: enable_encryption has been deprecated as encryption is automatically enabled if performing login with a username and password"

def enable_compression(client, enabled=1):
    """enable_compression has been deprecated and is here only for compatibility with old api.
    New client class will enable compression automatically and will use it if the send data will be smaller when compressed."""
    print "CTPRPC Warning: enable_compression has been deprecated as compression is automatically enabled and used"

def enable_data_compression(client, enabled=1):
    """enable_compression has been deprecated and is here only for compatibility with old api.
    New client class will enable compression automatically and will use it if the send data will be smaller when compressed."""
    print "CTPRPC Warning: enable_data_compression has been deprecated as compression is automatically enabled and used"

def event_subscribe(client, event_name, event_handler):
    """Subscribe the client to a remote event, when the server notifies the remote
    event to the client the event handler will be called"""
    client.event_subscribe(event_name, event_handler)

def event_unsubscribe(client, event_name, event_handler):
    """Unsubscribe the client from a remote event, the server will stop notifying the
    remote event to this client"""
    client.event_unsubscribe(event_name, event_handler)

def get_last_exception(client):
    """Last exception is already reraised when received from the remote method,
    so this function is probably useless and is here only for compatibility with old api"""
    return client.last_excp

def dump_last_exception(client):
    """Gets The last exception text output. Last remote exception traceback is
    already printed when raised by the remote end and so this function is probably useless
    and is here only for compatibility with old api"""
    return str(client.last_excp)

def print_last_exception(client):
    """Prints to screen the last remote exception traceback. Last remote exception traceback is
    already printed when raised by the remote end and so this function is probably useless
    and is here only for compatibility with old api"""
    print client.last_excp

def main():
    import sys, time

    host = 'localhost'
    port = 8000
    method = 'ping'
    call_args = []

    cur_opt = None
    for opt in sys.argv[1:]:
        if opt[0] == '-':
            cur_opt = opt
        else:
            if cur_opt == '-p':
                port = int(opt)
                cur_opt = None
            elif cur_opt == '-h':
                host = opt
                cur_opt = None
            elif cur_opt == '-m':
                method = opt
                cur_opt = None
            else:
                call_args.append(opt)

    cli = CTPRPCClient((host, port))
    call = getattr(cli, method)
    a = call(*call_args)
    print a

if __name__ == '__main__':
    main()
